// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <turbo/numeric/int128.h>

namespace nebula::compute::internal {

    // Accumulate sum/squared sum (using naive summation)
    // Shared implementation between scalar/hash aggregate variance/stddev kernels
    template<typename ArrowType>
    struct IntegerVarStd {
        using c_type = typename ArrowType::c_type;

        int64_t count = 0;
        int64_t sum = 0;
        turbo::int128 square_sum = 0;

        void ConsumeOne(const c_type value) {
            sum += value;
            square_sum += static_cast<uint64_t>(value) * value;
            count++;
        }

        double mean() const { return static_cast<double>(sum) / count; }

        double m2() const {
            // calculate m2 = square_sum - sum * sum / count
            // decompose `sum * sum / count` into integers and fractions
            const turbo::int128 sum_square = static_cast<turbo::int128>(sum) * sum;
            const turbo::int128 integers = sum_square / count;
            const double fractions = static_cast<double>(sum_square % count) / count;
            return static_cast<double>(square_sum - integers) - fractions;
        }
    };

    static inline void MergeVarStd(int64_t count1, double mean1, int64_t count2, double mean2,
                                   double m22, int64_t *out_count, double *out_mean,
                                   double *out_m2) {
        double mean = (mean1 * count1 + mean2 * count2) / (count1 + count2);
        *out_m2 += m22 + count1 * (mean1 - mean) * (mean1 - mean) +
                   count2 * (mean2 - mean) * (mean2 - mean);
        *out_count += count2;
        *out_mean = mean;
    }

}  // namespace nebula::compute::internal
